# -*- coding: utf-8 -*-
"""
Embedding客户端，处理文本向量化
"""

import asyncio
import json
import time
from typing import List, Dict, Any, Optional
import requests
from dataclasses import dataclass

from .config import embedding_config


@dataclass
class EmbeddingResult:
    """向量化结果"""
    vector: List[float]
    text: str
    model: str
    usage: Optional[Dict[str, Any]] = None


class EmbeddingClient:
    """Embedding客户端"""
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """初始化客户端"""
        if config is None:
            config = embedding_config.get_api_config()
        
        self.base_url = config["base_url"].rstrip("/")
        self.model = config["model"]
        self.api_key = config["api_key"]
        self.timeout = config["timeout"]
        self.max_retries = embedding_config.max_retries
        
        # 构建headers
        self.headers = {
            "Accept": "application/json",
            "Content-Type": "application/json",
        }
        if self.api_key:
            self.headers["Authorization"] = f"Bearer {self.api_key}"
    
    def _build_embedding_url(self) -> str:
        """构建embedding API URL"""
        if self.base_url.endswith("/v1/embeddings"):
            return self.base_url
        else:
            return f"{self.base_url}/v1/embeddings"
    
    def _make_request(self, texts: List[str], retry_count: int = 0) -> Dict[str, Any]:
        """发送embedding请求"""
        url = self._build_embedding_url()
        
        # 构建请求数据
        payload = {
            "model": self.model,
            "input": texts,
            "encoding_format": "float"
        }
        
        try:
            response = requests.post(
                url, 
                json=payload, 
                headers=self.headers, 
                timeout=self.timeout
            )
            response.raise_for_status()
            return response.json()
        except requests.HTTPError as e:
            if retry_count < self.max_retries:
                print(f"请求失败，重试第{retry_count + 1}次: {e}")
                time.sleep(2 ** retry_count)  # 指数退避
                return self._make_request(texts, retry_count + 1)
            else:
                raise Exception(f"Embedding请求失败: {e}")
        except Exception as e:
            raise Exception(f"Embedding请求异常: {e}")
    
    async def _make_request_async(self, texts: List[str], retry_count: int = 0) -> Dict[str, Any]:
        """异步发送embedding请求"""
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self._make_request, texts, retry_count)
    
    def embed_text(self, text: str) -> EmbeddingResult:
        """单个文本向量化"""
        if not text or not text.strip():
            raise ValueError("文本不能为空")
        
        print(f"开始向量化文本: {text[:50]}...")
        
        try:
            response = self._make_request([text])
            
            if "data" not in response or not response["data"]:
                raise Exception("API响应格式错误")
            
            embedding_data = response["data"][0]
            vector = embedding_data.get("embedding", [])
            
            if not vector:
                raise Exception("未获取到向量数据")
            
            result = EmbeddingResult(
                vector=vector,
                text=text,
                model=self.model,
                usage=response.get("usage")
            )
            
            print(f"向量化成功，向量维度: {len(vector)}")
            return result
            
        except Exception as e:
            print(f"向量化失败: {e}")
            raise
    
    async def embed_text_async(self, text: str) -> EmbeddingResult:
        """异步单个文本向量化"""
        if not text or not text.strip():
            raise ValueError("文本不能为空")
        
        print(f"开始异步向量化文本: {text[:50]}...")
        
        try:
            response = await self._make_request_async([text])
            
            if "data" not in response or not response["data"]:
                raise Exception("API响应格式错误")
            
            embedding_data = response["data"][0]
            vector = embedding_data.get("embedding", [])
            
            if not vector:
                raise Exception("未获取到向量数据")
            
            result = EmbeddingResult(
                vector=vector,
                text=text,
                model=self.model,
                usage=response.get("usage")
            )
            
            print(f"异步向量化成功，向量维度: {len(vector)}")
            return result
            
        except Exception as e:
            print(f"异步向量化失败: {e}")
            raise
    
    def embed_batch(self, texts: List[str]) -> List[EmbeddingResult]:
        """批量文本向量化"""
        if not texts:
            return []
        
        print(f"开始批量向量化，文本数量: {len(texts)}")
        
        results = []
        batch_size = embedding_config.batch_size
        
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i + batch_size]
            print(f"处理批次 {i//batch_size + 1}/{(len(texts) + batch_size - 1)//batch_size}")
            
            try:
                response = self._make_request(batch_texts)
                
                if "data" not in response or not response["data"]:
                    raise Exception("API响应格式错误")
                
                for j, embedding_data in enumerate(response["data"]):
                    vector = embedding_data.get("embedding", [])
                    if vector:
                        result = EmbeddingResult(
                            vector=vector,
                            text=batch_texts[j],
                            model=self.model,
                            usage=response.get("usage")
                        )
                        results.append(result)
                    else:
                        print(f"警告: 第{i+j}个文本向量化失败")
                        
            except Exception as e:
                print(f"批次 {i//batch_size + 1} 处理失败: {e}")
                # 继续处理下一批次
        
        print(f"批量向量化完成，成功处理 {len(results)}/{len(texts)} 个文本")
        return results
    
    async def embed_batch_async(self, texts: List[str]) -> List[EmbeddingResult]:
        """异步批量文本向量化"""
        if not texts:
            return []
        
        print(f"开始异步批量向量化，文本数量: {len(texts)}")
        
        results = []
        batch_size = embedding_config.batch_size
        
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i + batch_size]
            print(f"处理批次 {i//batch_size + 1}/{(len(texts) + batch_size - 1)//batch_size}")
            
            try:
                response = await self._make_request_async(batch_texts)
                
                if "data" not in response or not response["data"]:
                    raise Exception("API响应格式错误")
                
                for j, embedding_data in enumerate(response["data"]):
                    vector = embedding_data.get("embedding", [])
                    if vector:
                        result = EmbeddingResult(
                            vector=vector,
                            text=batch_texts[j],
                            model=self.model,
                            usage=response.get("usage")
                        )
                        results.append(result)
                    else:
                        print(f"警告: 第{i+j}个文本向量化失败")
                        
            except Exception as e:
                print(f"批次 {i//batch_size + 1} 处理失败: {e}")
                # 继续处理下一批次
        
        print(f"异步批量向量化完成，成功处理 {len(results)}/{len(texts)} 个文本")
        return results
